import re
import argparse
from selenium import webdriver
from selenium.common.exceptions import NoSuchElementException
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait
import time
from selenium.webdriver.common.keys import Keys
from loguru import logger
import shutil
import yaml
import os
from selenium.webdriver.common.action_chains import ActionChains
import json


def init_parser():
    parser = argparse.ArgumentParser(description="SQL2HiveSQL processing~")
    parser.add_argument("--fp", type=str, default="./sql", metavar="str",
                        help='file path of SQL file')
    parser.add_argument("--SQL", type=str, default="", metavar="str",
                        help='directly input single SQL statement')
    # parser.add_argument("--keyword",type=str, default="", metavar="str",
    #                     help='coral key words when searching Coral,using '|' to split every table. example:bytedance_opms_prd|bytedance_opms_prd|bytedance_opms_prd')
    parser.add_argument("--yaml", type=str, default="./conf/setting.yaml", metavar="str",
                        help='file path of setting-yaml file')
    parser.add_argument("--delimiter", type=str, default="empty_line", metavar="str",
                        help='extract SQL file into every single SQL statement by delimiter')
    parser.add_argument("--coral_key_words", type=str,
                        default="mds ods", metavar="str",
                        help="coral key words when searching Coral,using '|' to split every table")
    args_ = parser.parse_args()

    @logger.catch
    def init_setting(args_):
        global setting
        with open(args_.yaml, encoding='utf-8') as f:
            # f.read()此处读出来是字符串，用load方法转字典
            setting = yaml.load(stream=f.read(), Loader=yaml.FullLoader)
            print(setting)
        logger.info("============LOAD setting.yaml COMPLETE============")
        return setting

    setting_yaml = init_setting(args_)
    parser.add_argument("--regexp_rules", type=dict, default=setting_yaml['REGEXP_RULES'], metavar="dict",
                        help='regexp_rules,dict type,suggest edit the setting.yaml')
    parser.add_argument("--ua", type=str, default=setting_yaml['USER_AGENT'])
    parser.add_argument("--gufp", type=str, default=setting_yaml['GUFP'], metavar="str",
                        help='google user filepath, google user file which record cookies etc.')
    parser.add_argument("--gbfp", type=str, default=setting_yaml['GBFP'], metavar="str",
                        help='google executable file path')
    parser.add_argument("--egfp", type=str, default=setting_yaml['EGFP'], metavar="str",
                        help='EGToolKit filePath,copy google user filepath to EGFilePath,to avoid confict.')
    parser.add_argument("--account", type=str, default=setting_yaml['EMAIL'], metavar="str",
                        help='account to login SSO.')
    parser.add_argument("--pwd", type=str, default=setting_yaml['PWD'], metavar="str",
                        help='password to login SSO.')
    parser.add_argument("--driver_path", type=str, default=setting_yaml['DRIVER_PATH'], metavar="str",
                        help='chrome driver path')
    parser.add_argument("--time_wait", type=int, default=setting_yaml['TIMEWAIT'], metavar="int",
                        help='time wait (second) for each page loading')
    parser.add_argument("--coral_url", type=str, default=setting_yaml['CORAL_URL'], metavar="str")
    parser.add_argument("--aeolus_url", type=str, default=setting_yaml['AEOLUS_URL'], metavar="str")
    parser.add_argument("--aeolus_edit_xpath", type=str, default=setting_yaml['AEOLUS_EDIT_XPATH'], metavar="str")
    parser.add_argument("--cookie", type=str, default=setting_yaml['COOKIE'], metavar="str")
    global args
    args = parser.parse_args()
    assert args.fp or args.SQL,"必须指定SQL filePath 或者 输入 single SQL语句！"

def init_logger():
    logger.add(os.path.expanduser("./log/EGlog.log"))
    #logger.add('runtime_{time}.log', rotation='1 week', retention='10 week')

def init_driver():
    chrome_options = webdriver.ChromeOptions()
    #方便后台运行

    # chrome_options.add_argument('headless')
    # 设置开发者模式启动，该模式下webdriver属性为正常值
    chrome_options.add_experimental_option('excludeSwitches', ['enable-automation'])
    chrome_options.add_experimental_option('useAutomationExtension', False)
    #添加UA
    chrome_options.add_argument('user-agent='+args.ua)
    # 以最高权限运行
    chrome_options.add_argument('--no-sandbox')
     # 隐藏滚动条, 应对一些特殊页面
    chrome_options.add_argument('--hide-scrollbars')
    # 指定用户文件夹User Data路径
    chrome_options.add_argument(r'--user-data-dir='+args.gufp)
    # 指定缓存Cache路径
    chrome_options.add_argument(r"-–disk-cache-dir="+args.gufp+"/Cache")
    #设置chrome执行文件地址
    chrome_options.binary_location = args.gbfp

    global driver
    #Selenium执行cdp命令调用chrome浏览器的开发者工具，给window.navigator对象定义一个webdriver属性，并且设置为undefined,以绕过js的检测
    driver = webdriver.Chrome(executable_path=args.driver_path, options=chrome_options)
    script = '''
    Object.defineProperty(navigator, 'webdriver', {
        get: () => undefined
    })
    '''
    driver.execute_cdp_cmd("Page.addScriptToEvaluateOnNewDocument", {"source": script})

    #迁移Google浏览器用户目录防止打开窗口冲突
    if (os.path.exists(args.gufp)):
        shutil.rmtree(args.gufp)
    os.system('cp -r "{_gufp}" "{_egfp}"'.format(_gufp=args.gufp,_egfp=args.egfp))

def init_pages(driver):
    if not os.path.exists("./conf/cookies.json"):  # cookie未保存，保存
        driver.get(args.coral_url)
        while (not "https://data.bytedance.net/coral/" in driver.current_url):  # 登录进Coral，获得cookie开启aeolus
            print(driver.current_url)
            if (try_login_SSO(args.time_wait)): break
        # 保存当前登录的cookie
        dictCookies = driver.get_cookies()
        jsonCookies = json.dumps(dictCookies)
        # 登录完成后，将cookie保存到本地文件
        with open('./conf/cookies.json', 'w') as f:
            f.write(jsonCookies)
    else:
        with open('./conf/cookies.json', 'r', encoding='utf-8') as f:
            listCookies = json.loads(f.read())
        driver.get('https://data.bytedance.net/404') #表面domain头错误，需要绕开SSO访问页面，获取domain信息。惊讶的是字节的SSO居然没有防御404页面
        for cookie in listCookies:
            driver.add_cookie(cookie)
        driver.get(args.coral_url)
    global WINDOWS_HANDLES
    WINDOWS_HANDLES = driver.window_handles


def init():
    init_logger()
    init_parser()
    init_driver()
    init_pages(driver)

@logger.catch
def extract_sql():
    '''
    将多个SQL按照空行分割成单个的SQL列表
    :return:  分割好的SQL列表
    '''
    file = open(args.fp,'r').read().splitlines()
    single_state = []
    extract_res = []
    for line in file:
        if len(line)<1 and len(single_state)<1: #连续的空行或者开头空行
            continue
        elif len(line)<1: #将分割得到的single SQL加入结果集
            extract_res.append("\n".join(single_state))
            single_state.clear()
        else: #连续取数
            single_state.append(line)
    if len(single_state) > 0:extract_res.append("\n".join(single_state)) #最后剩一个表没空行
    logger.info("============EXTRACT SQL COMPLETE============")
    if args.coral_key_words:
        assert len(extract_res) == args.coral_key_words.count('|') + 1,"keywords个数必须和查询语句数目匹配！"
    return extract_res

@logger.catch
def SQL2HSQL(idx,single_state):
    '''
    将分割好的SQL列表经过多个正则规则转换为HiveSQL
    :param single_state: 分割好的SQL列表中的每个SQL
    :return: 转换好的HiveSQL
    '''

    #要求：输入partition 日期
    #选择selenium爬取出的表名

    son_depend_list = []
    table_names = []
    for regexp_rule in args.regexp_rules.keys():
        if regexp_rule not in son_depend_list:  # 不在依赖链中
            pattern = args.regexp_rules[regexp_rule]['pattern']
            replace = args.regexp_rules[regexp_rule]['replace'] if 'replace' in args.regexp_rules[
                regexp_rule] else None
            depends = args.regexp_rules[regexp_rule]['depends'] if 'depends' in args.regexp_rules[
                regexp_rule] else None
            if depends:  # 如果有子依赖，优先处理子依赖
                son_depend_list.extend(depends)  # 依赖链加入子依赖
                for depend in depends:
                    prog = re.compile(args.regexp_rules[depend]['pattern'])  # 子依赖的regexp
                    for father_pair in re.findall(pattern, single_state):  # 父依赖匹配
                        if prog.search(father_pair):  # 若匹配成功则在父依赖匹配结果中替换子依赖匹配结果
                            if 'find_table' in depend:  # 加入表名
                                table_names.extend(
                                    prog.findall(args.regexp_rules[depend]['pattern'], father_pair))
                            new_str = prog.sub(args.regexp_rules[depend]['replace'], father_pair)
                            single_state = single_state.replace(father_pair, new_str)
            elif replace:
                tns = re.findall(pattern, single_state)
                print(tns)
                if tns:
                    if 'find_table' in regexp_rule:  # 加入表名
                        table_names.extend(tns)
                    single_state = re.sub(pattern, replace, single_state)

    logger.info("FIND TABLE NUM : {_table_num} , TABLE NAMES ARE {_table_names}".format(_table_num=len(table_names),_table_names="   ".join(table_names)))
    print(single_state)
    #coral查出新表名并替换表名为hive分区的coral表
    keyword_list = args.coral_key_words.split('|')
    single_state_tables = list(set(table_names))
    print("single_state_tables:{_tns}".format(_tns=single_state_tables))
    for tn in single_state_tables:
        coral_tn = coral_replace(tn,tn+" "+keyword_list[idx])
        logger.info('Coral_Search_Replace_Table_Name {_origion_tn} to {_new_tn}'.format(_origion_tn=tn,_new_tn=coral_tn))
        #single_state = re.sub(r'\\b' + tn + '\\b', coral_tn, single_state) #replace未规定边界
        single_state = re.sub(tn, coral_tn, single_state)  # replace未规定边界
        print('statement: {_statement}'.format(_statement=single_state))

    #print('statement: {_statement}'.format(_statement=single_state))
    return single_state


def try_coral_search(keyword,time_wait):
    try:
        wait = WebDriverWait(driver, time_wait)  # 所有查询操作最大等待时间
        input_element = wait.until(EC.element_to_be_clickable(
            (By.XPATH,
             '//*[@id="app"]/section/main/div/div/div[2]/div[1]/div/div/div[2]/div/div/input')))
        if input_element:
            logger.info(r"=============进入Coral地图成功,搜索关键词{_keywords}================".format(_keywords=keyword))
            input_element.click()
            # 在输入框内输入值
            ActionChains(input_element).key_down(Keys.COMMAND).send_keys("a").key_up(Keys.COMMAND).send_keys(
                Keys.DELETE).perform()  # 输入框内容全部删除
            input_element.send_keys(Keys.ENTER)
            input_element.click()
            # 在输入框内输入值
            input_element.send_keys(keyword)
            # 回车
            input_element.send_keys(Keys.ENTER)
            return True
    except:
        logger.info("=============尝试寻找Coral地图失败================")
        return False

def try_login_SSO(time_wait):
    try:
        def is_login_success(time_wait):
            wait = WebDriverWait(driver, time_wait)  # 所有查询操作最大等待时间
            input_element = wait.until(EC.element_to_be_clickable(
                (By.XPATH,
                 '//*[@id="app"]/section/main/div/div/div[2]/div[1]/div/div/div[2]/div/div/input')))
            return input_element

        def is_enter_OTP(time_wait):
            wait = WebDriverWait(driver, time_wait)
            OTP_input = wait.until(EC.element_to_be_clickable(
                (By.XPATH,
                 '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/section/div/input')))
            return OTP_input

        def try_login_OTP(time_wait):
            try:
                wait = WebDriverWait(driver, time_wait)
                OTP_input = wait.until(EC.element_to_be_clickable(
                    (By.XPATH,
                     '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/section/div/input')))
                if OTP_input:
                    logger.info("=============进入OTP验证界面=================")
                    otp = input("请输入OTP密码  --   手机 SEAL 查询  ")
                    print("输入OTP："+otp)
                    OTP_input.click()
                    # OTP_input.clear()
                    OTP_input.send_keys(Keys.COMMAND,'a') #清除内容
                    OTP_input.send_keys(Keys.DELETE)
                    OTP_input.click()
                    OTP_input.send_keys(otp)
                    login_button = driver.find_element_by_xpath(
                        '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/section/button')
                    login_button.click()
                    logger.info("=============OTP登录中=================")
                    if (is_login_success(1)): return True  # 不需要OTP情况
                    else: return False
            except:
                logger.info("============OTP login timeout============")
                return False


        def try_login_pwd(time_wait):
            try:
                wait = WebDriverWait(driver, time_wait)  # 等待加载Coral最大时间，超时则说明需要SSO登录
                pwd_login_button = wait.until(EC.element_to_be_clickable(
                    (By.XPATH,
                     '//*[@id="root"]/div/section/div[1]/div/div[1]/div[2]')))
                if pwd_login_button:
                    pwd_login_button.click()

                login_account_input = wait.until(EC.element_to_be_clickable(
                    (By.XPATH,
                     '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/form/div[1]/div/div/input')))
                if login_account_input:
                    logger.info("=============进入密码登录界面=================")
                    login_account_input.click()
                    login_account_input.send_keys(args.account)
                    pwd = driver.find_element_by_xpath(
                        '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/form/div[2]/div/div/input')
                    pwd.click()
                    pwd.send_keys(args.pwd)
                    login_button = driver.find_element_by_xpath(
                        '//*[@id="root"]/div/section/div[1]/div/div[2]/div[2]/div/section/section/section/form/button')
                    login_button.click()
                if(is_enter_OTP(1)):return True
                else:return False
            except:
                logger.info("============login pwd timeout============")
                return False

        if not try_login_OTP(1): #不是第一次登录,停留OTP，需要先判断
            if try_login_pwd(args.time_wait):  #如果第一次就进入login流程
                return try_login_OTP(args.time_wait)  #login成功尝试OTP登录
        return False   #OTO校验登录成功

    except:
        logger.info("=============OSS验证失败================")
        return False

@logger.catch
def coral_replace(tn,keyword):
    '''
    在coral平台中将SQL表名替换为HiveSQL
    :param tn: 在SQL中搜索到的tableName
    :return: 查询好的tableName
    '''
    #1、搜索关键字
    #定位输入框元素
    while(not try_coral_search(keyword,1)):
        if(try_login_SSO(args.time_wait)):
            if(try_coral_search(keyword,args.time_wait)):break
    return coral_table_choice(keyword)


def coral_table_choice(keyword):
    # 2、获得coral反馈，由用户决定翻页
    logger.info("=============寻找Coral Table元素进行替换================")
    final_choice = 0
    page = 1
    full_page = None #记录最大page
    while(True):
        #1、搜定位coral返回信息，获得单页表名
        wait = WebDriverWait(driver, 15)#最大等待时间
        search_load = wait.until(EC.text_to_be_present_in_element_value( # input栏value为搜索词
            (By.XPATH,'//*[@id="app"]/section/main/div/div/div[2]/div[1]/div/div/div[2]/div/div/input'),keyword
        ))
        #url中包含搜索词
        former_page = wait.until(EC.element_to_be_clickable(
            (By.CLASS_NAME, 'byte-pagination-item')))
        search_complete = wait.until(EC.url_contains(keyword.replace(" ","%20")))
        table_name_elements = wait.until(EC.element_to_be_clickable((By.CLASS_NAME, 'name-link')))#等待前一页、后一页、表名加载完毕
        if search_load and search_complete and former_page and table_name_elements:
            pgs = driver.find_elements_by_class_name('byte-pagination-item') #获得所有页面
            if not full_page: #第一次  记录最大页数
                full_page = max([ele.text for ele in driver.find_elements_by_class_name("byte-pagination-item")])
            table_names = [element.text for element in driver.find_elements_by_class_name('name-link')]
            # print("element:"," ".join([element.text for element in elements]))

            for idx, tn in enumerate(table_names): #获得一页的所有表名
                logger.info(r"index of table is {_idx}, tableName is {_tn}".format(_idx=idx, _tn=tn))

            #2、提示用户选择，用户可以翻动页数
            time.sleep(0.1)
            input_string = input('数字选择表名，:q退出选择程序，输入a上一页，输入d下一页')
            logger.info(r"User Input |{_input}|".format(_input=input_string))
            if input_string == ":q":
                logger.info("~~强制退出~~")
                break
            elif input_string == 'a':
                if not page==1:page-=1
                logger.info(r"Action: jump to page {_page}".format(_page=page))
                print("pages:"+" ".join([pg.text for pg in pgs]))
                for pg in pgs:
                    if pg.text == page:
                        pg.click()
                        time.sleep(0.1)
                        continue
            elif input_string == 'd':
                if not page==full_page:page+=1
                logger.info(r"Action: jump to page {_page}".format(_page=page))
                print("pages:"+" ".join([pg.text for pg in pgs]))
                for pg in pgs:
                    if pg.text==page:
                        pg.click()
                        time.sleep(0.1)
                        continue
            elif '0' <= input_string <= '9':
                final_choice = table_names[int(input_string)]
                logger.info("~~正常退出~~")
                break
            else:
                logger.warning("~~请输入合法的字符!~~")
    return final_choice



def try_coral_search(keyword,time_wait):
    try:
        wait = WebDriverWait(driver, time_wait)  # 所有查询操作最大等待时间
        input_element = wait.until(EC.element_to_be_clickable(
            (By.XPATH,
             '//*[@id="app"]/section/main/div/div/div[2]/div[1]/div/div/div[2]/div/div/input')))
        if input_element:
            logger.info(r"=============进入Coral地图成功,搜索关键词{_keywords}================".format(_keywords=keyword))
            input_element.click()
            # 在输入框内输入值
            # input_element.clear()
            input_element.send_keys(Keys.COMMAND, 'a')  # 清除内容
            input_element.send_keys(Keys.DELETE)
            # 在输入框内输入值
            input_element.send_keys(keyword)
            # 回车
            input_element.send_keys(Keys.ENTER)
            time.sleep(0.1)
            return True
    except:
        logger.info("=============尝试寻找Coral地图失败================")
        return False

def exec_hsql(single_hsql,time_wait):
    f'''
    在Aeolus平台自动执行HSQL并下载
    :param single_hsql: 单个hsql
    :return: 执行结果 "hsql_idx" : "sucess"/"error statement"
    '''

    try:
        wait = WebDriverWait(driver, time_wait)  # 所有查询操作最大等待时间
        edit_aeolus = wait.until(EC.element_to_be_clickable(
            (By.CLASS_NAME,"inputarea")))
        submit_button = wait.until(EC.element_to_be_clickable(
            (By.XPATH,
             '//*[@id="guide-qe-run-task"]')))
        if edit_aeolus and submit_button:
            logger.info(r"=============进入Aeolus成功================")
            edit_aeolus.click()
            # 在输入框内输入值
            # edit_aeolus.clear()
            # edit_aeolus = driver.find_element_by_class_name('mtk1') #可交互的是小元素
            edit_aeolus.send_keys(Keys.COMMAND, 'a')  # 清除内容
            edit_aeolus.send_keys(Keys.DELETE)
            # 在输入框内输入值
            edit_aeolus.send_keys(single_hsql)
            ActionChains(driver).send_keys(Keys.SHIFT,Keys.ENTER).perform() #快捷键执行

            submit_button.click()
            error_occasion = driver.find_element_by_xpath('//*[@id="body"]/div[12]/div/div[2]/div/div[2]/div/div/div[1]/div/div/div[2]/button"]')

            error_occasion.click()
            logger.info("""
                                    ==============Aeolus 任务执行失败==============
                                    ==============     任务SQL      ==============
                                    {_state}
                                    """.format(_state=single_hsql))
            return False  # 执行出错
    except NoSuchElementException:
        logger.info("""
                                    ==============Aeolus 任务执行成功==============
                                    ==============     任务SQL      ==============
                                    {_state}
                                    """.format(_state=single_hsql))
        return True

    else:
        logger.info("""
                        ==============Aeolus 任务执行超时失败===========
                        ==============     任务SQL      ==============
                        {_state}
                        """.format(_state=single_hsql))
        return False





if __name__=="__main__":
    # try:
    # with ThreadPoolExecutor() as pool:
    # sqls = extract_sql()
    init()
    sqls = extract_sql()
    features = {}
    for idx,sql in enumerate(sqls):
        driver.switch_to.window(WINDOWS_HANDLES[0])
        statement = SQL2HSQL(idx,sql)
            # driver.switch_to.window(WINDOWS_HANDLES[1])
            # exec_hsql(statement,args.time_wait)
    # except:
    # finally:
    #     driver.quit()

